import ctypes
import sys
import threading
import time
import tkinter

import cv2
import keyboard
import numpy as np
import pyautogui
import pydirectinput

red_lower1 = np.array([0, 80, 236])
red_upper1 = np.array([10, 255, 255])
red_lower2 = np.array([150, 80, 236])
red_upper2 = np.array([190, 255, 255])

yellow_lower = np.array([20, 80, 236])
yellow_upper = np.array([30, 255, 255])

kernel = np.ones((3, 3), np.uint8)


def find_yellow_rectangles(hsv_img):
    mask = cv2.inRange(hsv_img, yellow_lower, yellow_upper)
    return find_rectangles(mask)


def find_red_rectangles(hsv_img):
    mask_red1 = cv2.inRange(hsv_img, red_lower1, red_upper1)
    mask_red2 = cv2.inRange(hsv_img, red_lower2, red_upper2)
    mask_red_combined = cv2.bitwise_or(mask_red1, mask_red2)
    return find_rectangles(mask_red_combined)


def find_rectangles(mask):
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if (w > 400 and h <= 30) or (w <= 30 and h > 300):
            return True
    return False


def is_admin():
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False


class App:
    def __init__(self):
        self.running = False
        self.root = tkinter.Tk()
        self.root.title("")
        self.root.geometry("176x144")
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)

        self.title_label = tkinter.Label(self.root, text="绝区零战斗副驾驶", bg="blue", fg="white")
        self.title_label.pack(pady=5)

        self.title_label = tkinter.Label(self.root, text="自动格挡和闪避")
        self.title_label.pack(pady=1)

        self.button_frame = tkinter.Frame(self.root)
        self.button_frame.pack(pady=1)

        self.start_button = tkinter.Button(self.button_frame, text="开始", command=self.start)
        self.start_button.pack(side=tkinter.LEFT, padx=5)

        self.stop_button = tkinter.Button(self.button_frame, text="停止", command=self.stop)
        self.stop_button.pack(side=tkinter.RIGHT, padx=5)

        self.status_label = tkinter.Label(self.root, text="状态：关闭", bg="red", fg="white")
        self.status_label.pack(pady=5)

        self.root.mainloop()

    def start(self):
        if not self.running:
            self.running = True
            self.status_label.config(text="状态：开启\n可短按 ` 键停止", bg="green")
            width, height = pyautogui.size()
            region = (0, 0, width, height)
            threading.Thread(target=self.auto_defense, args=(region,)).start()

    def stop(self):
        self.status_label.config(text="状态：关闭", bg="red")
        if self.running:
            self.running = False

    def on_closing(self):
        self.stop()
        self.root.destroy()

    def auto_defense(self, region):
        while self.running:
            if keyboard.is_pressed("`"):
                self.stop()
                break
            time.sleep(0.05)
            screenshot = pyautogui.screenshot(region=region)
            hsv_img = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2HSV)
            if find_yellow_rectangles(hsv_img):
                pydirectinput.press("space")
                time.sleep(0.3)
                continue
            if find_red_rectangles(hsv_img):
                pydirectinput.press("shift")
                time.sleep(0.3)


if __name__ == "__main__":
    if is_admin():
        app = App()

    else:
        ctypes.windll.shell32.ShellExecuteW(None, "runas", sys.executable, __file__, None, 1)
